package com.example.demo.security;

import com.example.demo.exception.JwtAuthenticationException;
import com.example.demo.model.entity.user.User;
import com.example.demo.model.enums.UserTypeEnum;
import io.jsonwebtoken.*;
import io.jsonwebtoken.security.Keys;
import io.jsonwebtoken.security.SignatureException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.authority.SimpleGrantedAuthority;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.crypto.SecretKey;
import java.nio.charset.StandardCharsets;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

@Slf4j
@Component
public class JwtTokenProvider {

    private static final String TOKEN_TYPE_CLAIM = "token_type";
    private static final String ROLES_CLAIM = "roles";
    private static final String USER_ID_CLAIM = "user_id";
    private static final String USER_TYPE_CLAIM = "user_type";
    private static final String ACCESS_TOKEN_TYPE = "access";
    private static final String REFRESH_TOKEN_TYPE = "refresh";
    private static final String ISSUER_CLAIM = "iss";
    private static final String AUDIENCE_CLAIM = "aud";

    private static final Set<String> revokedTokens = Collections.newSetFromMap(new ConcurrentHashMap<>());

    @Value("${app.jwt.secret}")
    private String jwtSecret;

    @Value("${app.jwt.access-token-expiration:1800000}")
    private long accessTokenExpirationMs;

    @Value("${app.jwt.refresh-token-expiration:604800000}")
    private long refreshTokenExpirationMs;

    @Value("${app.jwt.issuer:your-application-name}")
    private String jwtIssuer;

    @Value("${app.jwt.audience:your-audience}")
    private String jwtAudience;

    private SecretKey secretKey;

    @PostConstruct
    public void init() {
        validateConfig();
        this.secretKey = Keys.hmacShaKeyFor(jwtSecret.getBytes(StandardCharsets.UTF_8));
    }

    private void validateConfig() {
        if (jwtSecret == null || jwtSecret.trim().isEmpty()) {
            throw new IllegalStateException("JWT secret must not be null or empty");
        }
        if (jwtSecret.length() < 32) {
            log.warn("JWT secret is too short - consider using at least 32 characters for better security");
        }
        if (accessTokenExpirationMs <= 0) {
            throw new IllegalStateException("Access token expiration must be positive");
        }
        if (refreshTokenExpirationMs <= 0) {
            throw new IllegalStateException("Refresh token expiration must be positive");
        }
        if (refreshTokenExpirationMs <= accessTokenExpirationMs) {
            throw new IllegalStateException("Refresh token expiration must be longer than access token expiration");
        }
    }

    public String generateAccessToken(User user) {
        return generateToken(user, ACCESS_TOKEN_TYPE, accessTokenExpirationMs);
    }

    public String generateRefreshToken(User user) {
        return generateToken(user, REFRESH_TOKEN_TYPE, refreshTokenExpirationMs);
    }

    private String generateToken(User user, String tokenType, long expirationMs) {
        Instant now = Instant.now();
        Instant expiryDate = now.plus(expirationMs, ChronoUnit.MILLIS);

        List<String> roles = convertUserTypeToRoles(user.getUserType());

        return Jwts.builder()
                .setSubject(user.getUsername())
                .claim(USER_ID_CLAIM, user.getUserId())
                .claim(USER_TYPE_CLAIM, user.getUserType().name())
                .claim(TOKEN_TYPE_CLAIM, tokenType)
                .claim(ROLES_CLAIM, roles)
                .setIssuer(jwtIssuer)
                .setAudience(jwtAudience)
                .setIssuedAt(Date.from(now))
                .setExpiration(Date.from(expiryDate))
                .signWith(secretKey, SignatureAlgorithm.HS512)
                .compact();
    }

    private List<String> convertUserTypeToRoles(UserTypeEnum userType) {
        return Collections.singletonList(userType.name());
    }

    public String getUsernameFromToken(String token) {
        return validateAndGetClaims(token).getSubject();
    }

    public Long getUserIdFromToken(String token) {
        Claims claims = validateAndGetClaims(token);
        Object userId = claims.get(USER_ID_CLAIM);

        if (userId == null) {
            throw new JwtAuthenticationException("Token中缺少用户ID", "MISSING_USER_ID");
        }

        try {
            if (userId instanceof Number) {
                return ((Number) userId).longValue();
            }
            return Long.parseLong(userId.toString());
        } catch (NumberFormatException e) {
            throw new JwtAuthenticationException("无效的用户ID格式", "INVALID_USER_ID_FORMAT");
        }
    }

    public UserTypeEnum getUserTypeFromToken(String token) {
        try {
            String userType = validateAndGetClaims(token).get(USER_TYPE_CLAIM, String.class);
            return UserTypeEnum.valueOf(userType);
        } catch (IllegalArgumentException e) {
            throw new JwtAuthenticationException("无效的用户类型", "INVALID_USER_TYPE");
        }
    }

    public String getTokenTypeFromToken(String token) {
        return validateAndGetClaims(token).get(TOKEN_TYPE_CLAIM, String.class);
    }

    public List<GrantedAuthority> getAuthoritiesFromToken(String token) {
        try {
            List<String> roles = validateAndGetClaims(token).get(ROLES_CLAIM, List.class);
            return roles.stream()
                    .map(role -> role.startsWith("ROLE_") ? role : "ROLE_" + role)
                    .map(SimpleGrantedAuthority::new)
                    .collect(Collectors.toList());
        } catch (Exception e) {
            throw new JwtAuthenticationException("无法获取角色信息", "INVALID_ROLES_CLAIM");
        }
    }

    public boolean isAccessToken(String token) {
        return ACCESS_TOKEN_TYPE.equals(getTokenTypeFromToken(token));
    }

    public boolean isRefreshToken(String token) {
        return REFRESH_TOKEN_TYPE.equals(getTokenTypeFromToken(token));
    }

    public boolean validateToken(String token) {
        try {
            validateAndGetClaims(token);
            return true;
        } catch (JwtAuthenticationException e) {
            log.debug("Token validation failed: {}", e.getMessage());
            return false;
        }
    }

    public Claims validateAndGetClaims(String token) {
        if (isTokenRevoked(token)) {
            throw new JwtAuthenticationException("Token has been revoked", "JWT_REVOKED");
        }

        try {
            Claims claims = Jwts.parserBuilder()
                    .setSigningKey(secretKey)
                    .requireIssuer(jwtIssuer)
                    .requireAudience(jwtAudience)
                    .build()
                    .parseClaimsJws(token)
                    .getBody();

            // 验证必要声明
            if (claims.get(USER_ID_CLAIM) == null) {
                throw new JwtAuthenticationException("Missing user ID in token", "MISSING_USER_ID");
            }
            if (claims.get(USER_TYPE_CLAIM) == null) {
                throw new JwtAuthenticationException("Missing user type in token", "MISSING_USER_TYPE");
            }
            if (claims.get(TOKEN_TYPE_CLAIM) == null) {
                throw new JwtAuthenticationException("Missing token type in token", "MISSING_TOKEN_TYPE");
            }

            return claims;
        } catch (ExpiredJwtException ex) {
            throw new JwtAuthenticationException("Token expired", "JWT_EXPIRED");
        } catch (UnsupportedJwtException ex) {
            throw new JwtAuthenticationException("Unsupported JWT token", "JWT_UNSUPPORTED");
        } catch (MalformedJwtException ex) {
            throw new JwtAuthenticationException("Invalid JWT token", "JWT_INVALID");
        } catch (SignatureException ex) {
            throw new JwtAuthenticationException("JWT signature verification failed", "JWT_SIGNATURE_INVALID");
        } catch (IllegalArgumentException ex) {
            throw new JwtAuthenticationException("JWT claims string is empty", "JWT_CLAIMS_EMPTY");
        }
    }

    public void revokeToken(String token) {
        try {
            Claims claims = getClaimsFromToken(token);
            if (claims.getExpiration().after(new Date())) {
                revokedTokens.add(token);
                log.debug("Token revoked: {}", token);
            }
        } catch (JwtException e) {
            log.warn("Failed to revoke token: {}", e.getMessage());
        }
    }

    public void revokeTokensForUser(Long userId) {
        log.warn("revokeTokensForUser not fully implemented - consider adding token storage");
    }

    public boolean isTokenRevoked(String token) {
        return revokedTokens.contains(token);
    }

    private Claims getClaimsFromToken(String token) {
        return Jwts.parserBuilder()
                .setSigningKey(secretKey)
                .build()
                .parseClaimsJws(token)
                .getBody();
    }

    public long getAccessTokenExpirationMs() {
        return accessTokenExpirationMs;
    }

    public long getRefreshTokenExpirationMs() {
        return refreshTokenExpirationMs;
    }
}